"""
A4纸排版模块

用于将处理后的身份证照片排版到A4纸上
模拟扫描件效果，包括阴影、边框等细节
"""

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFilter
import math


class LayoutGenerator:
    """
    A4纸排版生成器
    
    功能：
    - 将图像放置在A4纸模拟背景上
    - 添加逼真的阴影效果
    - 模拟扫描件的视觉效果
    - 支持多种排版样式
    """
    
    def __init__(self):
        """初始化排版生成器"""
        # A4纸尺寸（300 DPI）
        self.a4_width = 2480   # 210mm * 300dpi / 25.4
        self.a4_height = 3508  # 297mm * 300dpi / 25.4
        
        # 身份证标准尺寸（300 DPI）
        self.id_width = 1012   # 85.6mm * 300dpi / 25.4  
        self.id_height = 638   # 54mm * 300dpi / 25.4
        
    def resize_to_id_card(self, image: np.ndarray) -> np.ndarray:
        """
        将图像调整为身份证标准尺寸
        
        参数:
            image: 输入图像数组
            
        返回:
            numpy.ndarray: 调整尺寸后的图像
            
        使用场景:
            generator = LayoutGenerator()
            resized = generator.resize_to_id_card(image_array)
        """
        height, width = image.shape[:2]
        
        # 计算缩放比例，保持宽高比
        scale_w = self.id_width / width
        scale_h = self.id_height / height
        scale = min(scale_w, scale_h)
        
        # 计算新尺寸
        new_width = int(width * scale)
        new_height = int(height * scale)
        
        # 调整大小
        resized = cv2.resize(image, (new_width, new_height), 
                           interpolation=cv2.INTER_LANCZOS4)
        
        # 如果尺寸不足，居中放置在白色背景上
        if new_width < self.id_width or new_height < self.id_height:
            result = np.ones((self.id_height, self.id_width, 3), 
                           dtype=np.uint8) * 255
            
            # 计算居中位置
            y_offset = (self.id_height - new_height) // 2
            x_offset = (self.id_width - new_width) // 2
            
            result[y_offset:y_offset+new_height, 
                   x_offset:x_offset+new_width] = resized
            
            return result
        
        return resized
    
    def create_shadow(self, image: np.ndarray, offset_x: int = 10, 
                     offset_y: int = 10, blur_radius: int = 15) -> np.ndarray:
        """
        为图像创建阴影效果
        
        参数:
            image: 输入图像
            offset_x: 阴影X偏移
            offset_y: 阴影Y偏移  
            blur_radius: 模糊半径
            
        返回:
            numpy.ndarray: 带阴影的图像
            
        使用场景:
            shadowed = generator.create_shadow(image, offset_x=8, offset_y=8)
        """
        # 转换为PIL图像以便处理阴影
        pil_image = Image.fromarray(image)
        
        # 创建阴影图层
        shadow = Image.new('RGBA', 
                          (pil_image.width + abs(offset_x) + blur_radius * 2,
                           pil_image.height + abs(offset_y) + blur_radius * 2),
                          (255, 255, 255, 0))
        
        # 创建阴影形状
        shadow_shape = Image.new('RGBA', pil_image.size, (0, 0, 0, 128))
        
        # 放置阴影
        shadow_x = blur_radius + max(0, offset_x)
        shadow_y = blur_radius + max(0, offset_y)
        shadow.paste(shadow_shape, (shadow_x, shadow_y))
        
        # 应用高斯模糊
        shadow = shadow.filter(ImageFilter.GaussianBlur(blur_radius))
        
        # 放置原图像
        img_x = blur_radius + max(0, -offset_x)
        img_y = blur_radius + max(0, -offset_y)
        shadow.paste(pil_image, (img_x, img_y))
        
        return np.array(shadow)
    
    def generate_a4_layout(self, image: np.ndarray, 
                          style: str = "center") -> np.ndarray:
        """
        生成A4纸排版布局
        
        参数:
            image: 处理后的身份证图像
            style: 排版样式 ("center", "top_left", "scan_style")
            
        返回:
            numpy.ndarray: A4纸排版后的图像
            
        使用场景:
            generator = LayoutGenerator()
            a4_image = generator.generate_a4_layout(id_image, "scan_style")
        """
        # 调整为身份证尺寸
        id_card = self.resize_to_id_card(image)
        
        # 创建A4背景（米白色，模拟纸张）
        a4_background = np.ones((self.a4_height, self.a4_width, 3), 
                               dtype=np.uint8) * 248
        
        if style == "center":
            # 居中放置
            y_pos = (self.a4_height - id_card.shape[0]) // 2
            x_pos = (self.a4_width - id_card.shape[1]) // 2
            
        elif style == "top_left":
            # 左上角放置（留边距）
            margin = 200
            y_pos = margin
            x_pos = margin
            
        elif style == "scan_style":
            # 模拟扫描件样式（稍微偏移，添加轻微旋转）
            y_pos = (self.a4_height - id_card.shape[0]) // 2 + 50
            x_pos = (self.a4_width - id_card.shape[1]) // 2 - 30
            
            # 添加轻微旋转（1-2度）
            angle = np.random.uniform(-1.5, 1.5)
            center = (id_card.shape[1] // 2, id_card.shape[0] // 2)
            rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
            id_card = cv2.warpAffine(id_card, rotation_matrix, 
                                   (id_card.shape[1], id_card.shape[0]),
                                   borderMode=cv2.BORDER_REFLECT)
        
        else:
            # 默认居中
            y_pos = (self.a4_height - id_card.shape[0]) // 2
            x_pos = (self.a4_width - id_card.shape[1]) // 2
        
        # 确保位置在有效范围内
        y_pos = max(0, min(y_pos, self.a4_height - id_card.shape[0]))
        x_pos = max(0, min(x_pos, self.a4_width - id_card.shape[1]))
        
        # 创建阴影效果
        shadowed_card = self.create_shadow(id_card)
        
        # 计算阴影放置位置
        shadow_y = max(0, y_pos - 15)
        shadow_x = max(0, x_pos - 15)
        
        # 确保阴影不会超出A4边界
        shadow_h, shadow_w = shadowed_card.shape[:2]
        if shadow_y + shadow_h > self.a4_height:
            shadow_h = self.a4_height - shadow_y
            shadowed_card = shadowed_card[:shadow_h, :]
        if shadow_x + shadow_w > self.a4_width:
            shadow_w = self.a4_width - shadow_x
            shadowed_card = shadowed_card[:, :shadow_w]
        
        # 处理透明度混合
        if shadowed_card.shape[2] == 4:  # RGBA
            alpha = shadowed_card[:, :, 3:] / 255.0
            rgb = shadowed_card[:, :, :3]
            
            # 混合到A4背景上
            for c in range(3):
                a4_background[shadow_y:shadow_y+shadow_h, 
                            shadow_x:shadow_x+shadow_w, c] = (
                    alpha.squeeze() * rgb[:, :, c] + 
                    (1 - alpha.squeeze()) * 
                    a4_background[shadow_y:shadow_y+shadow_h, 
                                shadow_x:shadow_x+shadow_w, c]
                )
        else:
            # 直接放置RGB图像
            a4_background[y_pos:y_pos+id_card.shape[0], 
                        x_pos:x_pos+id_card.shape[1]] = id_card
        
        return a4_background
    
    def add_scan_artifacts(self, image: np.ndarray) -> np.ndarray:
        """
        添加扫描仪效果（轻微噪声、压缩等）
        
        参数:
            image: 输入图像
            
        返回:
            numpy.ndarray: 添加扫描效果后的图像
            
        使用场景:
            scan_like = generator.add_scan_artifacts(a4_image)
        """
        # 添加轻微噪声
        noise = np.random.normal(0, 2, image.shape).astype(np.uint8)
        noisy = cv2.add(image, noise)
        
        # 轻微模糊模拟扫描效果
        blurred = cv2.GaussianBlur(noisy, (3, 3), 0.5)
        
        # 调整对比度模拟扫描仪特性
        alpha = 0.95  # 对比度因子
        beta = 5      # 亮度偏移
        adjusted = cv2.convertScaleAbs(blurred, alpha=alpha, beta=beta)
        
        return adjusted 